import ROOT
import numpy as np
import os
import csv

# 检查文件是否存在
file_path = "/opt/star-xp-master/build/source/0-Simulation/sim_4.root"
if not os.path.exists(file_path):
    print(f"文件 {file_path} 不存在，请检查路径。")
    import sys
    sys.exit(1)

# 加载 ROOT 字典文件
try:
    ROOT.gSystem.Load("../CommonLib/libSimEventDict.so")
except Exception as e:
    print(f"加载 ROOT 字典文件失败: {e}")
    import sys
    sys.exit(1)

# 打开 ROOT 文件
try:
    file = ROOT.TFile(file_path)
except Exception as e:
    print(f"打开 ROOT 文件失败: {e}")
    import sys
    sys.exit(1)

# 获取正确的树
tree = file.Get("Sim")
if not tree:
    print("未能找到 'Sim' 树，请检查文件内容。")
    file.Close()
    import sys
    sys.exit(1)

# 创建目录来存储 CSV 文件
output_dir = "extracted_data"
os.makedirs(output_dir, exist_ok=True)

# 遍历 ROOT 文件中的事件
for i in range(tree.GetEntries()):
    # 获取当前事件
    try:
        result = tree.GetEntry(i)
    except Exception as e:
        print(f"事件 {i}: 读取事件数据时出错: {e}")
        continue
    
    if result <= 0:
        print(f"事件 {i}: 无法读取事件数据，跳过该事件。")
        continue
    
    # 获取 SimEvent 对象
    sim_event = getattr(tree, "SimEvent", None)

    # 确认 sim_event 是否为 None
    if sim_event is None:
        print(f"事件 {i}: SimEvent 对象为空，跳过该事件。")
        continue

    # 提取 SimEvent 数据
    truth_energy = sim_event.GetTruthEnergy()
    truth_position = sim_event.GetTruthPosition()
    truth_polarization = sim_event.GetTruthPolarization()
    truth_particle = sim_event.GetTruthParticle()
    time_flag = sim_event.GetTimeFlag()

    track_map = sim_event.GetTrackMap()
    deposit_map = sim_event.GetDepositMap()

    # 初始化列表来存储数据
    tracks_info = []
    deposits_info = []

    # 处理 TrackMap
    if track_map and track_map.GetSize() > 0:
        track_iterator = track_map.MakeIterator()
        track_entry = track_iterator.Next()
        
        while track_entry:
            if isinstance(track_entry, ROOT.TObjString):
                key_tstring = track_entry.GetString()
                key_str = str(key_tstring)  # 将 TString 转换为 Python 字符串
                try:
                    value = track_map.GetValue(key_str)
                    if isinstance(value, ROOT.SimTrack):
                        pdgid = value.GetPDGID()
                        track_id = value.GetTrackID()
                        parent_id = value.GetParentID()
                        init_mass = value.GetInitMass()
                        init_ek = value.GetInitEk()
                        init_mom = value.GetInitMom()
                        init_pos = value.GetInitPos()
                        init_t = value.GetInitT()
                        exit_mom = value.GetExitMom()
                        exit_pos = value.GetExitPos()
                        exit_t = value.GetExitT()
                        track_length = value.GetTrackLength()
                        edep = value.GetEdep()
                        
                        step_idx = value.GetStepIdx()
                        steps = [str(idx) for idx in step_idx]

                        tracks_info.append({
                            "pdgid": pdgid,
                            "track_id": track_id,
                            "parent_id": parent_id,
                            "init_mass": init_mass,
                            "init_ek": init_ek,
                            "init_mom_x": init_mom.X(),
                            "init_mom_y": init_mom.Y(),
                            "init_mom_z": init_mom.Z(),
                            "init_pos_x": init_pos.X(),
                            "init_pos_y": init_pos.Y(),
                            "init_pos_z": init_pos.Z(),
                            "init_t": init_t,
                            "exit_mom_x": exit_mom.X(),
                            "exit_mom_y": exit_mom.Y(),
                            "exit_mom_z": exit_mom.Z(),
                            "exit_pos_x": exit_pos.X(),
                            "exit_pos_y": exit_pos.Y(),
                            "exit_pos_z": exit_pos.Z(),
                            "exit_t": exit_t,
                            "track_length": track_length,
                            "edep": edep,
                            "steps": ",".join(steps)
                        })
                except Exception as e:
                    print(f"事件 {i}, 获取轨道信息时出错: {e}")
            else:
                print(f"事件 {i}, 条目类型不匹配: 类型 {type(track_entry)}")

            track_entry = track_iterator.Next()

    # 处理 DepositMap
    if deposit_map and deposit_map.GetSize() > 0:
        deposit_iterator = deposit_map.MakeIterator()
        deposit_entry = deposit_iterator.Next()
        
        while deposit_entry:
            if isinstance(deposit_entry, ROOT.TObjString):
                key_tstring = deposit_entry.GetString()
                key_str = str(key_tstring)  # 将 TString 转换为 Python 字符串
                try:
                    value = deposit_map.GetValue(key_str)
                    if isinstance(value, ROOT.SimDeposit):
                        pdgid = value.GetPDGID()
                        track_id = value.GetTrackID()
                        parent_id = value.GetParentID()
                        charge = value.GetCharge()
                        pre_momentum = value.GetPreMomentum()
                        pre_position = value.GetPrePosition()
                        pre_t = value.GetPreT()
                        volume_name = value.GetVolumeName()
                        post_momentum = value.GetPostMomentum()
                        post_position = value.GetPostPosition()
                        post_t = value.GetPostT()
                        process_name = value.GetProcessName()
                        edep = value.GetEdep()
                        step_length = value.GetStepLength()
                        is_first_deposit = value.IsFirstDeposit()
                        
                        deposits_info.append({
                            "pdgid": pdgid,
                            "track_id": track_id,
                            "parent_id": parent_id,
                            "charge": charge,
                            "pre_mom_x": pre_momentum.X(),
                            "pre_mom_y": pre_momentum.Y(),
                            "pre_mom_z": pre_momentum.Z(),
                            "pre_pos_x": pre_position.X(),
                            "pre_pos_y": pre_position.Y(),
                            "pre_pos_z": pre_position.Z(),
                            "pre_t": pre_t,
                            "volume_name": volume_name.Data(),
                            "post_mom_x": post_momentum.X(),
                            "post_mom_y": post_momentum.Y(),
                            "post_mom_z": post_momentum.Z(),
                            "post_pos_x": post_position.X(),
                            "post_pos_y": post_position.Y(),
                            "post_pos_z": post_position.Z(),
                            "post_t": post_t,
                            "process_name": process_name.Data(),
                            "edep": edep,
                            "step_length": step_length,
                            "is_first_deposit": is_first_deposit
                        })
                except Exception as e:
                    print(f"事件 {i}, 获取沉积信息时出错: {e}")
            else:
                print(f"事件 {i}, 条目类型不匹配: 类型 {type(deposit_entry)}")

            deposit_entry = deposit_iterator.Next()

    # 如果有数据，则写入 CSV 文件
    output_file_path = os.path.join(output_dir, f"event_{i}.csv")
    with open(output_file_path, mode='w', newline='') as file_csv:
        writer = csv.writer(file_csv)
        writer.writerow([
            "truth_energy",
            "truth_position_x",
            "truth_position_y",
            "truth_position_z",
            "truth_polarization_s1",
            "truth_polarization_s2",
            "truth_polarization_s3",
            "truth_particle",
            "time_flag"
        ])
        writer.writerow([
            truth_energy,
            truth_position.X(),
            truth_position.Y(),
            truth_position.Z(),
            truth_polarization.Px(),  # Assuming S1 is stored in Px component
            truth_polarization.Py(),  # Assuming S2 is stored in Py component
            truth_polarization.Pz(),  # Assuming S3 is stored in Pz component
            truth_particle.Data(),
            time_flag
        ])

        if tracks_info:
            writer.writerow([])  # Empty row separator
            writer.writerow(["Tracks Info"])
            writer.writerow(tracks_info[0].keys())
            for track in tracks_info:
                writer.writerow(track.values())

        if deposits_info:
            writer.writerow([])  # Empty row separator
            writer.writerow(["Deposits Info"])
            writer.writerow(deposits_info[0].keys())
            for deposit in deposits_info:
                writer.writerow(deposit.values())

# 关闭 ROOT 文件
file.Close()
